// Copyright DustedPixels.com 2008. All rights reserved.

package com.dustedpixels.asm;

import java.util.ArrayList;
import java.util.HashMap;
import java.util.HashSet;
import java.util.List;
import java.util.Map;
import java.util.Set;

import org.objectweb.asm.MethodVisitor;
import org.objectweb.asm.Opcodes;
import org.objectweb.asm.Type;
import org.objectweb.asm.commons.AdviceAdapter;
import org.objectweb.asm.tree.AbstractInsnNode;
import org.objectweb.asm.tree.ClassNode;
import org.objectweb.asm.tree.FieldInsnNode;
import org.objectweb.asm.tree.FieldNode;
import org.objectweb.asm.tree.InsnNode;
import org.objectweb.asm.tree.MethodInsnNode;
import org.objectweb.asm.tree.MethodNode;
import org.objectweb.asm.tree.VarInsnNode;
import org.objectweb.asm.tree.analysis.Analyzer;
import org.objectweb.asm.tree.analysis.AnalyzerException;
import org.objectweb.asm.tree.analysis.Frame;
import org.objectweb.asm.tree.analysis.Interpreter;
import org.objectweb.asm.tree.analysis.SourceInterpreter;
import org.objectweb.asm.tree.analysis.SourceValue;
import org.objectweb.asm.tree.analysis.Value;


import com.dustedpixels.asm.analysis.Analyzers;
import com.dustedpixels.asm.analysis.FlowInterpreter;
import com.dustedpixels.asm.analysis.FlowValue;
import com.dustedpixels.asm.analysis.FrameUtils;
import com.dustedpixels.asm.analysis.LocalsInterpreter;
import com.dustedpixels.asm.analysis.TargetInterpreter;
import com.dustedpixels.asm.analysis.TargetValue;
import com.dustedpixels.asm.analysis.ThisInterpreter;
import com.dustedpixels.asm.utils.LocalsRemapper;
import com.dustedpixels.asm.utils.ReturnRemover;

/**
 * @author micapolos@gmail.com (Michal Pociecha-Los)
 */
public final class MethodInlinerUtils {
  public static void inlineInstance(ClassNode targetClass,
      ClassNode inlinedClass,
      FieldNode instanceField,
      String fieldNamePrefix) {
    
    removeField(targetClass, instanceField);
  }
  
  public static void removeField(ClassNode clazz, FieldNode field) {
    boolean removed = clazz.fields.remove(field);
    if (!removed) throw new IllegalArgumentException("No such field: " + field);
  }
  
  public static MethodNode inlineMethod(
      String ownerClassName,
      MethodNode targetMethod,
      MethodNode inlinedMethod) {
    
    MethodInsnNode methodInsn;
    
    do {
      methodInsn = findMethodInsn(ownerClassName, targetMethod, inlinedMethod);
      
      if (methodInsn != null) {
        targetMethod = inlineMethod(
            ownerClassName, 
            targetMethod, 
            methodInsn,
            inlinedMethod);
      }
    } while (methodInsn != null);
    
    return targetMethod;
  }

  public static MethodNode inlineMethod(
      String className,
      MethodNode method,
      MethodInsnNode methodInsn,
      MethodNode inlinedMethod) {

    // Validate descriptor of inlined method. It must match the descriptor of
    // method invocation instruction.
    if (!methodInsn.desc.equals(inlinedMethod.desc)) {
      throw new IllegalArgumentException("Method descriptor mismatch");
    }
    
    // Create a copy of target method for non-destructive processing.
    // Find corresponding method invocation instruction in copied method.
    int methodInsnIndex = method.instructions.indexOf(methodInsn);
    method = MethodNodes.copy(method);
    methodInsn = (MethodInsnNode) method.instructions.get(methodInsnIndex);

    // Analyze for local variables flow.
    int maxLocal = AnalysisUtils.calculateMaxLocal(
        className,
        method, 
        methodInsn);
    
    // Analyze for local variables flow.
    int stackSize = AnalysisUtils.calculateStackSize(
        className,
        method, 
        methodInsn);
    
    // Process method to be inlined
    MethodNode processedInlinedMethod = MethodNodes.shallowCopy(inlinedMethod);
    MethodVisitor processor = LocalsRemapper.shift(processedInlinedMethod, maxLocal);
    processor = ReturnRemover.newOn(processor);
    inlinedMethod.accept(processor);
    inlinedMethod = processedInlinedMethod;
    
    // Store method arguments in locals.
    Type[] argTypes = MethodNodes.getAllArgTypes(inlinedMethod);
    int[] argLocals = MethodNodes.allocateLocals(argTypes, maxLocal);
    for (int i = argTypes.length - 1; i >= 0; i--) {
      Type argType = argTypes[i];
      int argLocal = argLocals[i];
      System.out.println("Arg " + i + " at " + argLocal);
      AbstractInsnNode storeInsn =
          new VarInsnNode(OpcodeUtils.getStoreOpcode(argType), argLocal);
      method.instructions.insertBefore(methodInsn, storeInsn);
    }

    // Inline bytecode
    method.instructions.insertBefore(methodInsn, inlinedMethod.instructions);
    
    // Remove method-call instruction.
    method.instructions.remove(methodInsn);
            
    // Update maxs
    method.maxLocals = 
        Math.max(
            method.maxLocals,
            argLocals[argTypes.length] + inlinedMethod.maxLocals); 
    method.maxStack = 
      Math.max(
          method.maxStack,
          stackSize + inlinedMethod.maxStack); 
    
    return method;
  }
  
  /**
   * Renumbers local variable indices in given method to start from given value.
   * 
   * @param method a method
   * @param startLocal start index of locals
   * @return a method with renumbered locals
   */
  @SuppressWarnings("unchecked")
  public static MethodNode renumberLocals(MethodNode method, final int startLocal) {
    MethodNode resultMethod = new MethodNode(
        method.access, 
        method.name, 
        method.desc,
        method.signature, 
        (String[]) method.exceptions.toArray(new String[0]));
    
    MethodVisitor remapper = LocalsRemapper.shift(resultMethod, startLocal);
    
    method.accept(remapper);
    
    return resultMethod;
  }
  
  /**
   * Removes all RETURN instructions from given method replacing them with GOTO
   * to the end of the method (if necessary).
   * 
   * @return processed method
   */
  @SuppressWarnings("unchecked")
  public static MethodNode replaceReturnWithGoto(MethodNode method) {
    String[] exceptions = new String[method.exceptions.size()];
    method.exceptions.toArray(exceptions);
    MethodNode resultMethod = new MethodNode(
        method.access, 
        method.name, 
        method.desc,
        method.signature, 
        exceptions);
    
    MethodVisitor returnGotoReplacer = ReturnRemover.newOn(resultMethod);
    
    method.accept(returnGotoReplacer);

    return resultMethod;
  }
  
  /**
   * Find an instruction in given method, which calls other method on "this".
   * @param method
   * @return
   */
  public static MethodInsnNode findMethodInsn(
      String ownerClassName,
      MethodNode targetMethod,
      MethodNode invokedMethod) {
    
    // Analyze method code looking for references to "this".
    Interpreter thisInterpreter = new ThisInterpreter();
    Analyzer thisAnalyzer = new Analyzer(thisInterpreter);
    Frame[] thisFrames;
    try {
      thisFrames = thisAnalyzer.analyze(ownerClassName, targetMethod);
    } catch (AnalyzerException e) {
      throw new RuntimeException(e);
    }
    
    AbstractInsnNode[] instructions = targetMethod.instructions.toArray();
    
    for (int i = 0; i < instructions.length; i++) {
      AbstractInsnNode instruction = instructions[i];
      if (instruction instanceof MethodInsnNode) {
        MethodInsnNode methodInsn = (MethodInsnNode) instruction;
        Frame thisFrame = thisFrames[i];
        
        if (methodInsn.name.equals(invokedMethod.name) &&
            methodInsn.desc.equals(invokedMethod.desc)) {
          /*
           * At this point we know that method name and arguments match. Now,
           * we need to check whether this method is invoked on 'this' object.
           */
          int argCount = Type.getArgumentTypes(invokedMethod.desc).length;
          Value thisValue = thisFrame.getStack(thisFrame.getStackSize() - argCount - 1);
          if (thisValue == ThisInterpreter.Values.THIS) {
            return methodInsn;
          }
        }
      }
    }
    
    return null;
  }
  
  /**
   * Returns a map mapping inlined method argument index onto target method
   * local index, which can be reused.
   * 
   * @param ownerClassName
   * @param targetMethod
   * @param targetMethodInsn
   * @param inlinedMethod
   * @return
   */
  public static Map<Integer, Integer> calculateInlineArgReuse(
      String ownerClassName,
      MethodNode targetMethod,
      MethodInsnNode targetMethodInsn,
      MethodNode inlinedMethod) {
    
    try {
      // Calculate the number of stack entries which stores method arguments,
      // including reference to "this" for non-static methods.
      Type[] argTypes = MethodNodes.getAllArgTypes(inlinedMethod);
  
      // Initialize an array of stack reuse.
      Map<Integer, Integer> reuse = new HashMap<Integer, Integer>();
      
      // Get index of method invocation instruction in target method bytecode.
      int methodInsnIndex = targetMethod.instructions.indexOf(targetMethodInsn);
      
      // Analyze for local variables flow.
      FlowInterpreter valueFlowInterpreter = new FlowInterpreter();
      Analyzer valueFlowAnalyzer = Analyzers.newUsing(valueFlowInterpreter);
      Frame[] valueFlowFrames = valueFlowAnalyzer.analyze(ownerClassName, targetMethod);
      Frame valueFlowFrame = valueFlowFrames[methodInsnIndex];
      
      LocalsInterpreter localsInterpreter = new LocalsInterpreter();
      Analyzer localsAnalyzer = Analyzers.newUsing(localsInterpreter);
      Frame[] localsFrames = localsAnalyzer.analyze(ownerClassName, targetMethod);
      Frame localsFrame = localsFrames[methodInsnIndex];
      
      Set<Integer> writtenLocals = MethodNodes.findWrittenLocals(targetMethod);
      
      int localIndex = 0;
      int argIndex = 0; 
      int stackIndex = localsFrame.getStackSize() - argTypes.length - 1;
      for (int i = 0; i < argTypes.length; i++) {
        if (!writtenLocals.contains(localIndex)) {
          int targetLocal = LocalsInterpreter.getVar(localsFrame.getStack(stackIndex));
          if (targetLocal != -1) {
            if (valueFlowFrame.getStack(stackIndex) == valueFlowFrame.getLocal(targetLocal)) {
              reuse.put(argIndex, targetLocal);
            }
          }
        }
        Type argType = argTypes[i];
        localIndex += argType.getSize();
        stackIndex--;
        argIndex++;
      } 

      return reuse;
    } catch (AnalyzerException e) {
      throw new RuntimeException(e);
    }
  }
  
  public static void removeUnusedLoadStoreInsns(String className, MethodNode method) {
    try {
      Interpreter sourceInterpreter = new SourceInterpreter();
      Interpreter targetInterpreter = new TargetInterpreter();
      Interpreter flowInterpreter = new FlowInterpreter();
      
      Analyzer sourceAnalyzer = Analyzers.newUsing(sourceInterpreter);
      Analyzer targetAnalyzer = Analyzers.newUsing(targetInterpreter);
      Analyzer flowAnalyzer = Analyzers.newUsing(flowInterpreter);
      
      Frame[] sourceFrames = sourceAnalyzer.analyze(className, method);
      Frame[] targetFrames = targetAnalyzer.analyze(className, method);
      Frame[] flowFrames = flowAnalyzer.analyze(className, method);
      
      AbstractInsnNode[] insns = method.instructions.toArray();
      List<AbstractInsnNode> insnsToRemove = new ArrayList<AbstractInsnNode>();
      
      for (int i = 0; i < insns.length; i++) {
        AbstractInsnNode insn = insns[i];
        if (OpcodeUtils.isStore(insn.getOpcode())) {          
          int var = ((VarInsnNode) insn).var;
          System.out.println("Got ISTORE " + var + " at " + i);
          
          Frame sourceFrame = sourceFrames[i];
          Frame flowFrame = flowFrames[i];
          Frame nextTargetFrame = targetFrames[i + 1];          
          Frame nextFlowFrame = flowFrames[i + 1];
          
          Value flowValue = nextFlowFrame.getLocal(var);

          Set<AbstractInsnNode> candidateInsnsToRemove = new HashSet<AbstractInsnNode>();
          boolean ok = false;
          TargetValue targetValue = (TargetValue) nextTargetFrame.getLocal(var);
          Set<AbstractInsnNode> targets = targetValue.getTargets();
          if (targets.isEmpty()) {
            
          } else {
            ok = true;
            for (AbstractInsnNode targetInsn : targetValue.getTargets()) {
              if (!OpcodeUtils.isLoad(targetInsn.getOpcode())) {
                System.out.println(" Found target which is not LOAD - dismissing");
                ok = false;
                break;
              } else {
                int targetInsnIndex = method.instructions.indexOf(targetInsn);
                Frame targetFlowFrame = flowFrames[targetInsnIndex];
                if (targetFlowFrame.getLocal(var) == flowValue) {
                  System.out.println(" Found LOAD " + var + " at " + targetInsnIndex + " accessing the same local");
                  candidateInsnsToRemove.add(targetInsn);
                } else {
                  System.out.println(" Found LOAD " + var + " at " + targetInsnIndex + " accessing different local - dismissing");
                  ok = false;
                  break;
                }
              }
            }
          }
                
          if (ok) {
            // Add instructions to remove
            insnsToRemove.addAll(candidateInsnsToRemove);
            insnsToRemove.add(insn);
          }
        }
      }
      
      for (AbstractInsnNode insnToRemove : insnsToRemove) {
        method.instructions.remove(insnToRemove);
      }
    } catch (AnalyzerException e) {
      throw new RuntimeException(e);
    }
  }
  
  public static MethodNode replaceFieldWithLocal(
      final String className,
      MethodNode method) {
    // Perform analysis on copy.
    method = MethodNodes.copy(method);
    
    try {
      Interpreter flowInterpreter = new FlowInterpreter();
      Interpreter sourceInterpreter = new SourceInterpreter();
      Interpreter targetInterpreter = new TargetInterpreter();
      
      Analyzer flowAnalyzer = Analyzers.newUsing(flowInterpreter);
      Analyzer sourceAnalyzer = Analyzers.newUsing(sourceInterpreter);
      Analyzer targetAnalyzer = Analyzers.newUsing(targetInterpreter);
      
      Frame[] flowFrames = flowAnalyzer.analyze(className, method);
      Frame[] sourceFrames = sourceAnalyzer.analyze(className, method);
      Frame[] targetFrames = targetAnalyzer.analyze(className, method);
      
      AbstractInsnNode[] insns = method.instructions.toArray();

      // Collect all GETFIELD and PUTFIELD instructions, which 
      final Set<FieldInsnNode> fieldInsns = new HashSet<FieldInsnNode>();
      final Map<String, String> fieldName2DescMap = new HashMap<String, String>();
      final Set<String> uncertainFieldNames = new HashSet<String>();
      
      for (int i = 0; i < insns.length; i++) {
        AbstractInsnNode insn = insns[i];
        if (insn instanceof FieldInsnNode) {
          FieldInsnNode fieldInsn = (FieldInsnNode) insn;
          if (insn.getOpcode() == Opcodes.GETFIELD || 
              insn.getOpcode() == Opcodes.PUTFIELD) {
            int stackShift = (insn.getOpcode() == Opcodes.GETFIELD) ? 0 : 1;
            Frame flowFrame = flowFrames[i];
            FlowValue flowValue = (FlowValue) FrameUtils.top(flowFrame, stackShift);
            FlowValue thisValue = (FlowValue) flowFrames[0].getLocal(0);
            if (flowValue == thisValue) {
              //System.out.println("Field access at: " + i);
              fieldInsns.add(fieldInsn);
              fieldName2DescMap.put(fieldInsn.name, fieldInsn.desc);
            } else if (flowValue.covers(thisValue)) {
              //System.out.println("Uncertain field access at: " + i);
              uncertainFieldNames.add(fieldInsn.name);
            }
          }
        }
      }
      
      // Replace field instructions with load/store.
      final Set<AbstractInsnNode> insnsToRemove = new HashSet<AbstractInsnNode>();
      final Set<FieldInsnNode> fieldInsnsWithPop = new HashSet<FieldInsnNode>();
      final Set<String> fieldNames = new HashSet<String>();
      final Set<String> getFieldNames = new HashSet<String>();
      final Set<String> putFieldNames = new HashSet<String>();
      
      for (FieldInsnNode fieldInsn : fieldInsns) {
        if (!uncertainFieldNames.contains(fieldInsn.name)) {
          fieldNames.add(fieldInsn.name);
          if (fieldInsn.getOpcode() == Opcodes.GETFIELD) {
            getFieldNames.add(fieldInsn.name);
          } else {
            putFieldNames.add(fieldInsn.name);
          }
          //System.out.println("Processing field " + fieldInsn.name);
          int fieldInsnIndex = method.instructions.indexOf(fieldInsn);
          int stackShift = (fieldInsn.getOpcode() == Opcodes.GETFIELD) ? 0 : 1;
          Frame sourceFrame = sourceFrames[fieldInsnIndex];
          SourceValue sourceValue = (SourceValue) FrameUtils.top(sourceFrame, stackShift);
          AbstractInsnNode sourceInsnToRemove = null;
          //System.out.println(" Number of sources: " + sourceValue.insns.size());
          if (sourceValue.insns.size() == 1) {
            AbstractInsnNode sourceInsn = (AbstractInsnNode) sourceValue.insns.iterator().next();
            //System.out.println(" Source instruction: " + sourceInsn);
            if (sourceInsn.getOpcode() == Opcodes.ALOAD ||
                sourceInsn.getOpcode() == Opcodes.DUP) {
              int sourceInsnIndex = method.instructions.indexOf(sourceInsn);
              TargetValue targetValue = (TargetValue) FrameUtils.top(targetFrames[sourceInsnIndex + 1]);
              //System.out.println(" Number of targets: " + targetValue.getTargets().size());
              if (targetValue.getTargets().size() == 1) {
                sourceInsnToRemove = sourceInsn;
              }
            }
          }
  
          if (sourceInsnToRemove == null) {
            fieldInsnsWithPop.add(fieldInsn);
          } else {
            insnsToRemove.add(sourceInsnToRemove);
          }
        }
      }
      
      // Generate locals for fields
      int nextLocal = method.maxLocals;
      final Map<String, Integer> fieldName2LocalMap = new HashMap<String, Integer>();
      for (String fieldName : fieldNames) {
        fieldName2LocalMap.put(fieldName, nextLocal);
        String fieldDesc = fieldName2DescMap.get(fieldName);
        Type fieldType = Type.getType(fieldDesc);
        int fieldSize = fieldType.getSize();
        nextLocal += fieldSize;
      }
      
      // Remove instructions.
      for (AbstractInsnNode insnToRemove : insnsToRemove) {
        method.instructions.remove(insnToRemove);
      }
      
      for (FieldInsnNode fieldInsn : fieldInsns) {
        VarInsnNode varInsn;
        if (fieldInsn.getOpcode() == Opcodes.GETFIELD) {
          varInsn = new VarInsnNode(
              OpcodeUtils.getLoadOpcode(Type.getType(fieldInsn.desc)), 
              fieldName2LocalMap.get(fieldInsn.name));
        } else {
          varInsn = new VarInsnNode(
              OpcodeUtils.getStoreOpcode(Type.getType(fieldInsn.desc)),
              fieldName2LocalMap.get(fieldInsn.name));
        }

        if (fieldInsnsWithPop.contains(fieldInsn)) {
          System.out.println(" Inserting POP");
          InsnNode popInsn = new InsnNode(Opcodes.POP);
          method.instructions.insertBefore(fieldInsn, popInsn);
        }
        method.instructions.insertBefore(fieldInsn, varInsn);
        method.instructions.remove(fieldInsn);
      }
      
      // Wrap method with field->local and local->field transfer
      MethodNode newMethod = MethodNodes.shallowCopy(method);
      MethodVisitor wrapper = new AdviceAdapter(newMethod, method.access, method.name, method.desc) {
        @Override protected void onMethodEnter() {
          for (String fieldName : getFieldNames) {
            String fieldDesc = fieldName2DescMap.get(fieldName);
            super.visitVarInsn(Opcodes.ALOAD, 0);
            super.visitFieldInsn(Opcodes.GETFIELD, 
                className,
                fieldName, 
                fieldName2DescMap.get(fieldName));
            super.visitVarInsn(
                OpcodeUtils.getStoreOpcode(Type.getType(fieldDesc)),
                fieldName2LocalMap.get(fieldName));
          }
        }

        @Override
        protected void onMethodExit(int opcode) {
          for (String fieldName : putFieldNames) {
            String fieldDesc = fieldName2DescMap.get(fieldName);
            super.visitVarInsn(Opcodes.ALOAD, 0);
            super.visitVarInsn(
                OpcodeUtils.getLoadOpcode(Type.getType(fieldDesc)),
                fieldName2LocalMap.get(fieldName));
            super.visitFieldInsn(Opcodes.PUTFIELD, 
                className,
                fieldName, 
                fieldName2DescMap.get(fieldName));
          }
        }
      };
      
      method.instructions.resetLabels();
      method.accept(wrapper);
      
      newMethod.maxLocals = nextLocal;
      
      return newMethod;
    } catch (AnalyzerException e) {
      throw new RuntimeException(e);
    }
  }
}
